package com.raos.example;

import com.raos.example.service.OllamaService;
import lombok.extern.slf4j.Slf4j;
import org.noear.solon.annotation.Inject;
import org.noear.solon.net.annotation.ServerEndpoint;
import org.noear.solon.net.websocket.WebSocket;
import org.noear.solon.net.websocket.WebSocketListener;

import java.io.IOException;
import java.nio.ByteBuffer;
import java.util.concurrent.CompletableFuture;

/**
 * WebSocket服务 处理器
 *
 * @author raos
 * @date 2025-02-14
 */
@Slf4j
@ServerEndpoint("/ws/chat/{user}")
public class WebSocketHandler implements WebSocketListener {

    @Inject("${ollama.chat.model}")
    private String modelName;

    @Inject
    private OllamaService ollamaService;

    @Override
    public void onOpen(WebSocket socket) {
        log.info("开启连接，会话id=[{}]...", socket.id());
        // path var
        String user = socket.param("user");
        // query var
        String token = socket.param("token");
        if (!("admin".equals(user) && "123456".equals(token))) {
            log.info("[{}]账号秘钥错误，关闭连接...", user);
            socket.send("账号秘钥错误，即将退出会话...");
            socket.close();
        }
    }

    @Override
    public void onMessage(WebSocket socket, String text) throws IOException {
        log.info("会话id=[{}]，接收消息：[{}]", socket.id(),text);
        // 异步处理请求
        CompletableFuture.runAsync(() -> {
            try {
                final StringBuffer line = new StringBuffer();
                ollamaService.generateResponse(modelName, text, response -> {
                    log.debug("异步响应内容：{}", response);
                    // socket.send(response);
                    line.append(response);
                    if (response.endsWith("\n")) {
                        socket.send(line.toString());
                        line.setLength(0);
                    }
                    else if ("".equals(response)) { // 最后输出空串代表结束
                        socket.send(line.toString());
                    }
                });
            } catch (Exception e) {
                socket.send("服务暂时不可用");
            }
        });
    }

    @Override
    public void onMessage(WebSocket socket, ByteBuffer binary) throws IOException {
        onMessage(socket, binary.toString());
    }

    @Override
    public void onClose(WebSocket socket) {
        log.info("会话id=[{}]，关闭连接...", socket.id());
    }

    @Override
    public void onError(WebSocket socket, Throwable error) {
        log.error("会话id=[{}]，连接异常...", socket.id(), error);
    }

}
